import cv2
import numpy as np
import sys
import onnxruntime as ort
import yaml


def preprocess(img, size):
    img_w, img_h = img.shape[1], img.shape[0]
    w, h = size
    new_w = int(img_w * min(w * 1. / img_w, h * 1. / img_h))
    new_h = int(img_h * min(w * 1. / img_w, h * 1. / img_h))
    resized_image = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    canvas = np.full((h, w, 3), 0)
    canvas[(h - new_h) // 2:(h - new_h) // 2 + new_h, (w - new_w) // 2:(w - new_w) // 2 + new_w, :] = resized_image
    input_array = canvas.transpose([2, 0, 1]) / 255.
    return input_array[np.newaxis,:].astype(np.float32)




def main():
    configs = yaml.load(open('./config.yaml').read(),Loader=yaml.FullLoader)
    print(1)
    model = ort.InferenceSession(configs['model'])
    cap = cv2.VideoCapture(configs['rtsp'])
    assert cap.isOpened(), 'Cannot capture source'
    while cap.isOpened():
        ret, img_ori = cap.read()
        if ret:
            img = preprocess(img_ori, configs['input_size'])
            outputs_onnx = model.run(None, {'input': img})[0]
            pred = np.sum(outputs_onnx/100)
            nums = int(pred+0.5)
            level = 'comfortable'
            for key, value in configs['level'].items():
                if len(value) == 1:
                    if nums >= value[0]:
                        level = key
                    else:
                        continue
                elif len(value) == 2:
                    if nums >= value[0] and nums <= value[1]:
                        level = key
                    else:
                        continue
                else:
                    raise Exception("Invalid value!", value)
            text = f'nums: {nums}  level: {level}'
            img_ori = cv2.putText(img_ori,text,(50,50),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1)
            cv2.imshow('ori',img_ori)
            if cv2.waitKey(1) == 27:
                break


if __name__ == "__main__":
    main()

